import re
from argparse import Namespace

import torch
import torch.nn.functional as F

from PytorchNLP3Yelp.ReviewClassifier import ReviewClassifier
from PytorchNLP3Yelp.ReviewDataset import ReviewDataset


def preprocess_text(text):
    text=text.lower()
    text=re.sub(r"([.,!?])",r" \1",text)
    text=re.sub(r"[^a-zA-Z.,!?]+",r" ",text)
    return text

def predict_rating(review,classifier,vectorizer,decision_threshold=0.5):
    review=preprocess_text(review)
    vectorized_review=torch.tensor(vectorizer.vectorize(review))
    result=classifier(vectorized_review.view(1,-1))

    probability_value=torch.sigmoid(result).item()
    index=1
    if probability_value<decision_threshold:
        index=0

    return vectorizer.rating_vocab.lookup_index(index)

if __name__=='__main__':
    test_review="this is a pretty awesome book"

    args = Namespace(
        frequency_cutoff=25,
        model_state_file='model.pth',
        review_csv='data/yelp/reviews_with_splits_lite.csv',
        save_dir='model/yelp/',
        vectorizer_file='vectorizer.json',
        batch_size=128,
        early_stopping_criteria=5,
        learning_rate=0.001,
        num_epochs=100,
        seed=1337,
        catch_keyboard_interrupt=True,
        cuda=True,
        expand_filepaths_to_save_dir=True,
        reload_from_files=False,
    )

    if not torch.cuda.is_available():
        args.cuda = False
    print("Using CUDA: {}".format(args.cuda))
    args.device = torch.device("cuda" if args.cuda else "cpu")

    if args.reload_from_files:
        # 如果有词向量文件
        print("Loading dataset and vectorizer")
        dataset = ReviewDataset.load_dataset_and_load_vectorizer(args.review_csv,
                                                                 args.vectorizer_file)
    else:
        # 如果没有词向量文件，创建
        print("Loading dataset and creating vectorizer")
        dataset = ReviewDataset.load_dataset_and_make_vectorizer(args.review_csv)
        dataset.save_vectorizer(args.vectorizer_file)

    vectorizer = dataset.get_vectorizer()
    classifier = ReviewClassifier(num_features=len(vectorizer.review_vocab))
    classifier=classifier.cpu()
    prediction=predict_rating(test_review,classifier,vectorizer,decision_threshold=0.5)
    print("{}->{}".format(test_review,prediction))

    print(classifier.fc1.weight.shape)

    fc1_weights=classifier.fc1.weight.detach()[0]
    _,indices=torch.sort(fc1_weights,dim=0,descending=True)
    indices=indices.numpy().tolist()

    print("积极的前20个词")
    for i in range(20):
        print(vectorizer.review_vocab.lookup_index(indices[i]))

    print("===\n\n\n")

    print("消极的前20个词")
    indices.reverse()
    for i in range(20):
        print(vectorizer.review_vocab.lookup_index(indices[i]))